
import uproot
import pandas as pd
sig_name = {"../../Data/data_4tau/signal/THDM_hA_4tau_mh_mch_ATLAS_pt10_14TeV_01.root": 3.67719}
bkg_name = {
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbar.root": 4452.73,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarww.root": 0.0139,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarz.root": 0.2136,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarzz.root": 0.0001,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_wwjj.root": 268.291,
    #"/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zjj.root": 151060,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zz.root": 9.296
}

df_sig_list = []
for name, value in sig_name.items():
    df_tmp = uproot.open(name)["datatrain"].arrays(["h1_inm", "h2_inm"], library="pd")
    df_tmp["weight"] = value/len(df_tmp)
    print(name, value/len(df_tmp))
    df_sig_list.append(df_tmp)

df_bkg_list = []
for name, value in bkg_name.items():
    df_tmp = uproot.open(name)["datatrain"].arrays(["h1_inm", "h2_inm"], library="pd")
    df_tmp["weight"] = value/len(df_tmp)
    print(name, value/len(df_tmp))
    df_bkg_list.append(df_tmp)

# print(file.keys()) >>> "datatrain"
# print(tree.keys()) >>> "pt", "met" ...
# here are ['weight', 'jet_num', 'signal_jet_num', 'jet_pt1', 'jet_eta1', 'jet_pt2', 'jet_eta2', 'elec_num', 'muon_num', 'lep_num', 'signal_lep_num', 'lep_pt1', 'lep_eta1', 'lep_pt2', 'lep_eta2', 'photon_num', 'MET', 'vis_m', 'vis_pt', 'vis_eta', 'll_inm', 'll_pt', 'll_eta', 'jj_inm', 'jj_pt', 'jj_eta', 'h1_inm', 'h1_pt', 'h1_eta', 'h2_inm', 'h2_pt', 'h2_eta', 'H_inm', 'H_pt', 'H_eta', 'l_l_openangle', 'j_j_openangle', 'l1_j1_openangle', 'l1_j12_openangle', 'l2_j1_openangle', 'l2_j2_openangle', 'h1_h2_openangle', 'HT', 'samesign', 'MT2']

# df = tree.arrays(["h1_inm", "h2_inm"],library="pd")

# df = df[(df!= -10000).all(axis=1)]
df_sig = pd.concat(df_sig_list)
df_bkg = pd.concat(df_bkg_list)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 确定直方图的边界和 bin 数量
bins = 40
dz = 0
for df in df_bkg_list:
    # 在计算直方图时考虑权重
    hist, xedges, yedges = np.histogram2d(df['h1_inm'], df['h2_inm'], bins=bins, range=[[0, 200], [0, 200]], weights=df["weight"])

    # 生成网格坐标
    xpos, ypos = np.meshgrid(xedges[:-1], yedges[:-1])
    xpos = xpos + (xedges[1] - xedges[0]) / 2
    ypos = ypos + (yedges[1] - yedges[0]) / 2
    zpos = np.zeros_like(xpos)
    dx = xedges[1] - xedges[0]
    dy = yedges[1] - yedges[0]
    dz1 = hist
    # 将 dz 转换为与 xpos 和 ypos 匹配的形状
    dz = dz + dz1.T.flatten()

# 创建 3D 图形
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 绘制 3D 直方图
ax.bar3d(xpos.flatten(), ypos.flatten(), zpos.flatten(), dx, dy, dz, alpha=0.6, zsort='average', color="red")

# 在计算信号直方图时考虑权重
hist, xedges, yedges = np.histogram2d(df_sig['h1_inm'], df_sig['h2_inm'], bins=bins, range=[[0, 200], [0, 200]], weights=df_sig["weight"])
dz2 = hist
# 将 dz 转换为与 xpos 和 ypos 匹配的形状
dz2 = dz2.T.flatten()
#ax.bar3d(xpos.flatten(), ypos.flatten(), zpos.flatten() + dz, dx, dy, dz2, alpha=0.5, zsort='average', color="blue")
ax.bar3d(xpos.flatten(), ypos.flatten(), zpos.flatten() , dx, dy, dz2, alpha=0.5, zsort='average', color="blue")

# 设置标题和坐标轴标签
ax.set_xlabel(r'$M_{lj}^{1}$')
ax.set_ylabel(r'$M_{lj}^{2}$')
ax.set_zlabel('Event Number')

plt.savefig('mh1_mh2_hist.pdf')
# 显示图形
plt.show()
